#include "FreeJoint.h"
#include "RigidBodyNode.h"
#include "Mat6x6.h"
#include "PrescribedMotion.h"
#include "EOMData.h"
#include "ABAData.h"

using namespace RSIM;
using namespace std;

///////////////////////////////////////////////////////
// Free Joint
///////////////////////////////////////////////////////

void Free::printJointType() const{cout<<"Joint Type=Free\n";}

///////////////////////////////////////////////////////

int Free::getJointDOF() const{return 6;}

///////////////////////////////////////////////////////

void Free::setDefaultQ(double *q)const{
        q[0] = 1.0;	q[1] = 0.0;
        q[2] = 0.0;	q[3] = 0.0;
        q[4] = 0.0;	q[5] = 0.0;
        q[6] = 0.0;
}

///////////////////////////////////////////////////////

void Free::setDefaultU(double *u)const{
        u[0] = 0.0;	u[1] = 0.0;
        u[2] = 0.0;	u[3] = 0.0;
        u[4] = 0.0;	u[5] = 0.0;
}

///////////////////////////////////////////////////////

Free::Free(const int& NodeIndex):Joint(FreeJoint){
        this->NodeIndex_ = NodeIndex;
}

///////////////////////////////////////////////////////

void Free::updJointTransform(	const double *q,
                                        const RigidBodyNode *Parent,
                                        const int& pChildID,
                                        double *betak){
        // Q(0:3) Euler parameters
        this->T_MF_.wgetR().setRotationFromQuaternion(q);
        
        // Q(4:6) v_MF
        this->T_MF_.wgetP().copy(q+4);		
        
        // update betak
        this->updBetaK(q, Parent, pChildID, betak);
}

///////////////////////////////////////////////////////

void Free::updBetaK(	const double *q,
                        const RigidBodyNode *Parent,
                        const int& pChildID,
                        double *betak)const{
        if(Parent){
                const double *T_FM_P_Ptr = Parent->data_->getBodyTransformFM()[pChildID].P().Ptr();
                const double *T_FM_R_Ptr = Parent->data_->getBodyTransformFM()[pChildID].R().Ptr();
                
                FastOps::mult_mat33_vec3(T_FM_R_Ptr,q+4,betak);

                betak[0] += T_FM_P_Ptr[0];
                betak[1] += T_FM_P_Ptr[1];
                betak[2] += T_FM_P_Ptr[2];
        }
}

///////////////////////////////////////////////////////

Free* Free::New(const int& NodeIndex){
        return new Free(NodeIndex);
}

///////////////////////////////////////////////////////

void Free::updVelocity(	const State& state, 
                        const RigidBodyNode *Parent,
                        const Rotation& prF_C_F, 
                        const double *betak,
                        Vect3& w,
                        Vect3& v) const {
                                        
        const Vect3& pr_w = Parent->data_->getAngularVelocity();
        const Vect3& pr_v = Parent->data_->getLinearVelocity();
        
        // 1. F_C_prF*pr_w is the absolute angular velocity
        // of the parent expressed in the joint's F frame
        // 2. Vect3(U[0],U[1],U[2]) is the angular velocity of this body
        // which is already expressed in the joint's M frame
        
        double *W = w.wPtr();
        {
                double tmp[3];
                FastOps::transpose_mult_mat33_vec3(prF_C_F.Ptr(), pr_w.Ptr(), W);
                
                FastOps::transpose_mult_mat33_vec3(this->T_MF_.R().Ptr(),
                                                state.getUForNode(this->NodeIndex_),
                                                tmp);
                W[0] += tmp[0];
                W[1] += tmp[1];
                W[2] += tmp[2];
        }
        double v_tmp[3], rel_tmp[3];
        
        // omega x r
        FastOps::cross(pr_w.Ptr(),betak,v_tmp);
        
        // v + omega x r
        v_tmp[0] += pr_v(0);
        v_tmp[1] += pr_v(1);
        v_tmp[2] += pr_v(2);
        
        double *v_wptr = v.wPtr();
        
        // convert v + omega x r to F basis of this body
        FastOps::transpose_mult_mat33_vec3(prF_C_F.Ptr(), v_tmp, v_wptr);
        
        FastOps::transpose_mult_mat33_vec3(this->T_MF_.R().Ptr(),
                                        state.getUForNode(this->NodeIndex_)+3,
                                        rel_tmp);
        
        // Velocity = v + omega x r + qdot
        v_wptr[0] += rel_tmp[0];
        v_wptr[1] += rel_tmp[1];
        v_wptr[2] += rel_tmp[2];
}

///////////////////////////////////////////////////////

void Free::updRelVelocity(const State& state,Vect3& w, Vect3& v) const {
        
        // exressed in F basis of this body
        FastOps::transpose_mult_mat33_vec3(T_MF_.R().Ptr(),
                                        state.getUForNode(this->NodeIndex_),
                                        w.wPtr());
                                        
        FastOps::transpose_mult_mat33_vec3(T_MF_.R().Ptr(),
                                        state.getUForNode(this->NodeIndex_)+3,
                                        v.wPtr());
}

///////////////////////////////////////////////////////

void Free::calc_H_x_UDot(const State& state,double *h_x_udot) const{	
        FastOps::transpose_spatial_basis_shift_vec(this->T_MF_.R().Ptr(),
                                                state.getUDotForNode(this->NodeIndex_),
                                                h_x_udot);
}

///////////////////////////////////////////////////////

// Equation 2.188 \cite{Farid}
void Free::calcQDot(State& state)const{
        double *QDot = state.wgetQDotForNode(this->NodeIndex_);
        const double *U = state.getUForNode(this->NodeIndex_);
        const double *Q = state.getQForNode(this->NodeIndex_);
        
        QDot[0] = 0.5*(-Q[1]*U[0] - Q[2]*U[1] - Q[3]*U[2]);
        QDot[1] = 0.5*( Q[0]*U[0] + Q[3]*U[1] - Q[2]*U[2]);
        QDot[2] = 0.5*(-Q[3]*U[0] + Q[0]*U[1] + Q[1]*U[2]);
        QDot[3] = 0.5*( Q[2]*U[0] - Q[1]*U[1] + Q[0]*U[2]);
        
        QDot[4] = U[3];
        QDot[5] = U[4];
        QDot[6] = U[5];
}

///////////////////////////////////////////////////////

void Free::calcAKT(	const RigidBodyNode *Parent, 
                        const RigidBodyNode *Body,
                        const int& ChildID,
                        const double *betak,
                        State& state,
                        Vect6& sAkt) const{

        double *UDot = state.wgetUDotForNode(this->NodeIndex_);

        // AKT is calculated in F basis of this body
        if(Parent){
                const double *U = state.getUForNode(this->NodeIndex_);

                double *akt_ptr = sAkt.wPtr();
                const double *prFCF_Ptr = Body->data_->getRotationFCF().Ptr();

                // Angular velocity of parent expressed in F of parent
                const Vect3& parent_ang = Parent->data_->getAngularVelocity();

                {
                        double wwr[3];
                        // cross( w^{k-1}, cross(w^{k-1}, \beta^{k-1}) ) {expressed in F of parent}
                        FastOps::wxwxr(parent_ang.Ptr(), betak, wwr);
                        
                        // cross( w^{k-1}, cross(w^{k-1}, \beta^{k-1}) ) {expressed in F of this body}
                        FastOps::transpose_mult_mat33_vec3(prFCF_Ptr, wwr, akt_ptr+3);			
                }

                double wkk_x_u[3],ang_vel[3];

                // Covert Angular velocity of parent to F basis of this body
                FastOps::transpose_mult_mat33_vec3(prFCF_Ptr, parent_ang.Ptr(),ang_vel);

                // convert w^k_k u^k to F basis of the parent
                FastOps::transpose_mult_mat33_vec3(this->T_MF_.R().Ptr(), U, wkk_x_u);
                FastOps::cross(ang_vel, wkk_x_u,akt_ptr);

                double wxvkk[3],s_U[3];
                FastOps::transpose_mult_mat33_vec3(this->T_MF_.R().Ptr(), U+3,s_U);
                
                // 2*cross(w^k, v^{k+1}_{k+1}u^k+1)
                FastOps::cross(ang_vel,s_U,wxvkk);
                akt_ptr[3] += 2*wxvkk[0];
                akt_ptr[4] += 2*wxvkk[1];
                akt_ptr[5] += 2*wxvkk[2];
                
                // In this case, Akt = Akt + p^k_k \dot{u}^k
                if(this->isMotionPrescribedAndActive_){
                                
                        // Update QDot_ cache from prescribed motion
                        this->PM_->setUDot(state.getSysData(),UDot);
                                
                        // convert w^k_k dot(u)^kto F basis of parent
                        double wkk_x_udot[6];
                
                        // Akt in this case has the additional term wkk_x_udot which is already known
                        FastOps::transpose_spatial_basis_shift_vec(this->T_MF_.R().Ptr(), UDot,wkk_x_udot);
                        
                        akt_ptr[0] += wkk_x_udot[0]; 
                        akt_ptr[1] += wkk_x_udot[1]; 
                        akt_ptr[2] += wkk_x_udot[2];
                        akt_ptr[3] += wkk_x_udot[3]; 
                        akt_ptr[4] += wkk_x_udot[4]; 
                        akt_ptr[5] += wkk_x_udot[5];
                }
        }
        else{ 
                if(this->isMotionPrescribedAndActive_){
                        this->PM_->setUDot(state.getSysData(),UDot);
                }
        }
}

///////////////////////////////////////////////////////

void Free::sweepBackwards(const bool& isLeaf, 
                const int& ChildID,
                const Vect6& sAkt,
                const Rotation& FCF,					    
                const EOMData *eom,
                const double *betak,
                RigidBodyNode *Parent,
                ABAData *aba) const{		
        
        // get writable pointers for articulated body inertia of this body.
        // Ik3[0] is the ABI of this body and the eqns of motion for this body
        // can be written as Ik3[0] * A - Fk3[0] = 0. Here the ABI of intermediate frames 
        // are calculated and stored in Ik3+1 and Ik3+2
        Mat6x6 *Ik3 = aba->wgetIk3();
        Vect6 *Fk3 = aba->wgetFk3();	
        
        if(!isMotionPrescribedAndActive_){
                
                this->calc_Triang_x_Ik3_TransZ(Ik3->Ptr(),(Ik3+1)->wPtr());
                this->calc_Triang_x_cFk_TransZ(Ik3->Ptr(), Fk3->Ptr(), (Fk3+1)->wPtr());		

                this->calc_Triang_x_Ik3_TransY((Ik3+1)->Ptr(), (Ik3+2)->wPtr());	
                this->calc_Triang_x_cFk_TransY((Ik3+1)->Ptr(), (Fk3+1)->Ptr(), (Fk3+2)->wPtr());
                
                this->calc_Triang_x_Ik3_TransX((Ik3+2)->Ptr(), (Ik3+3)->wPtr());	
                this->calc_Triang_x_cFk_TransX((Ik3+2)->Ptr(), (Fk3+2)->Ptr(), (Fk3+3)->wPtr());
                
                this->calc_Triang_x_Ik3_RotZ((Ik3+3)->Ptr(),(Ik3+4)->wPtr());
                this->calc_Triang_x_cFk_RotZ((Ik3+3)->Ptr(), (Fk3+3)->Ptr(), (Fk3+4)->wPtr());		

                this->calc_Triang_x_Ik3_RotY((Ik3+4)->Ptr(), (Ik3+5)->wPtr());	
                this->calc_Triang_x_cFk_RotY((Ik3+4)->Ptr(), (Fk3+4)->Ptr(), (Fk3+5)->wPtr());

                if(Parent){
                        this->calc_cFk((Fk3+5)->Ptr(),(Ik3+5)->Ptr(),sAkt.Ptr(),aba->wgetcFk()->wPtr());			
                }
        } // if (!isMotionPrescribedAndActive)
        else{
                if(Parent){
                        // cFk = Fk3 - Ik3*A^{k-1}_t
                        this->calc_cFk(Fk3->Ptr(),Ik3->Ptr(),sAkt.Ptr(),aba->wgetcFk()->wPtr());
                }
        }
}

/////////////////////////////////////////////////////////////////////////

void Free::sweepForward(	  const RigidBodyNode *Parent,
                                const int& ChildID,
                                const Vect6& sAkt,
                                const Rotation& FCF,
                                const ABAData *aba,
                                const double *betak,
                                EOMData *eom,
                                State& state
                                ){
        
        this->calcQDot(state);	
        
        const double *R_MF_Ptr = this->T_MF_.R().Ptr();	
        const Vect6 *Fk3 = aba->getFk3();
        const Mat6x6 *Ik3 = aba->getIk3();
        Vect6 *Ak1 = eom->wgetAk1();
        double *ak1_ptr = Ak1->wPtr();
        double *UDot = state.wgetUDotForNode(this->NodeIndex_);
        
        if(Parent){
                const double *Parent_Ak1_Ptr = Parent->data_->getEOMData()->getAk1()->Ptr();
                if(!isMotionPrescribedAndActive_){
                        double udot_tmp[6];
                        {
                                double ak1_tmp[6];
                                FastOps::transpose_skbeta_x_V(Parent_Ak1_Ptr, betak, ak1_tmp);
                                FastOps::transpose_spatial_basis_shift_vec(FCF.Ptr(),ak1_tmp, ak1_ptr);
                                this->calc_UDot_RotX((Ik3+5)->Ptr(),ak1_ptr,aba->getcFk()->Ptr(),udot_tmp);
                        }
                        Ak1->add(sAkt.Ptr());
                        ak1_ptr[0] += udot_tmp[0];

                        this->calc_UDot_RotY((Ik3+4)->Ptr(),ak1_ptr,(Fk3+4)->Ptr(),udot_tmp+1);
                        ak1_ptr[1] += udot_tmp[1];
                        
                        this->calc_UDot_RotZ((Ik3+3)->Ptr(),ak1_ptr,(Fk3+3)->Ptr(),udot_tmp+2);
                        ak1_ptr[2] += udot_tmp[2];
                        
                        this->calc_UDot_TransX((Ik3+2)->Ptr(),ak1_ptr,(Fk3+2)->Ptr(),udot_tmp+3);
                        ak1_ptr[3] += udot_tmp[3];
                        
                        this->calc_UDot_TransY((Ik3+1)->Ptr(),ak1_ptr,(Fk3+1)->Ptr(),udot_tmp+4);
                        ak1_ptr[4] += udot_tmp[4];
                        
                        this->calc_UDot_TransZ(Ik3->Ptr(),ak1_ptr,Fk3->Ptr(),udot_tmp+5);
                        ak1_ptr[5] += udot_tmp[5];
                        
                        // udot_tmp is UDot expressed in F basis of this body.
                        // Must be converted back to M basis of parent. 
                        FastOps::spatial_basis_shift_vec(R_MF_Ptr, udot_tmp, UDot);

                }//if(!isMotionPrescribedAndActive_)
                else{
                        double Ak1p1[6];
                        // Ak1p1 (0:5)= transpose(skbeta)*A^{parent}_1 
                        // In this case, sAkt already includes p^k_k \dot{u}^k
                        FastOps::transpose_skbeta_x_V(Parent_Ak1_Ptr, betak,Ak1p1);
                        
                        // Ak1 = Ak1p1 (0:5) exprssed in F of this body
                        FastOps::transpose_spatial_basis_shift_vec(FCF.Ptr(),Ak1p1,eom->wgetAk1()->wPtr());
                        eom->wgetAk1()->add(sAkt.Ptr());
                }
        }//if(Parent)
        else{
                if(!isMotionPrescribedAndActive_){	
                        const double *fk3_ptr = Fk3[5].Ptr();
                        ak1_ptr[0] = fk3_ptr[0]/Ik3[5].Ptr()[0];
                        ak1_ptr[1]=0.0; ak1_ptr[2]=0.0; ak1_ptr[3]=0.0;
                        ak1_ptr[4]=0.0; ak1_ptr[5]=0.0;
                        
                        this->calc_UDot_RotY((Ik3+4)->Ptr(),ak1_ptr,(Fk3+4)->Ptr(),ak1_ptr+1);
                        this->calc_UDot_RotZ((Ik3+3)->Ptr(),ak1_ptr,(Fk3+3)->Ptr(),ak1_ptr+2);
                        this->calc_UDot_TransX((Ik3+2)->Ptr(),ak1_ptr,(Fk3+2)->Ptr(),ak1_ptr+3);
                        this->calc_UDot_TransY((Ik3+1)->Ptr(),ak1_ptr,(Fk3+1)->Ptr(),ak1_ptr+4);
                        this->calc_UDot_TransZ(Ik3->Ptr(),ak1_ptr,Fk3->Ptr(),ak1_ptr+5);

                        FastOps::spatial_basis_shift_vec(R_MF_Ptr, ak1_ptr, UDot);
                }
                else{
                        // Calculate Ak1_ in F frame of this body
                        double *ak1_ptr = eom->wgetAk1()->wPtr();
                        FastOps::spatial_basis_shift_vec(R_MF_Ptr, UDot, ak1_ptr);
                }		
        }
        // calculate the spatial constraint force
        this->calc_m_cFk(aba->getFk3()[0].Ptr(), aba->getIk3()[0].Ptr(), Ak1->Ptr(), eom->wgetFkc1()->wPtr());
}